"""
@Time: 2020/12/11 下午 6:20
@Author: jinzhuan
@File: fn_predictor_test.py
@Desc: 
"""
import torch
from cognlp import *


torch.cuda.set_device(4)
device = torch.device('cuda')
vocabulary = Vocabulary.load('../data/fn/framenet/data/vocabulary.txt')
model = Bert4Fn(len(vocabulary), device=device)
predictor = FnPredictor(model, model_path='../data/fn/framenet/model/2020-12-11-15:36:31/checkpoint-5960/model.pt',
                        vocabulary=vocabulary, device=device, device_ids=[4, 5],
                        lu_path="../data/fn/framenet/data/lu_key_mapping.pkl")
print(predictor.predict("located in the eastern Pyrenees mountains and bordered by Spain and France. Created under a charter in A.D. 988, the present Principality was formed in A.D. 1278. It is known as a principality as it is a monarchy headed by two Co-Princes – the Spanish/Roman Catholic Bishop of Urgell and the President of France."))
